#!/usr/bin/python3
# -*- coding: utf-8 -*-
import os
import struct
import multiprocessing

video_list_count = 8
# base_path = r'/rootfs/media/kasim/DataSet/Record'
# out_path = r'/rootfs/media/kasim/DataSet/Record'
base_path = r'/rootfs/media/kasim/DataSet/RecordEx'
out_path = r'/rootfs/media/kasim/DataSet/RecordEx'
video_list_bbox_count_file = r'video_list{}_bbox_count.txt'
video_list_bbox_file = r'video_list{}_bbox.dat'
video_list_filter_file = r'ad_filter_low_select_file_list.txt'

PROCESS_COUNT = video_list_count

select_class = 'Face'

select_class_label = {
    'Person': 0,
    'Cat': 1,
    'Dog': 2,
    'BabyCar': 3,
    'Face': 4,
}
select_label = select_class_label[select_class]

THRESHOLDS = [
    0.3,  # Person
    0.3,  # Cat
    0.3,  # Dog
    0.3,  # BabyCar
    0.3,  # Face
]


def filter_proc(file_queue, id):
    try:
        filter_video_dict = set()
        with open(os.path.join(base_path, video_list_filter_file), 'r') as file:
            for video_info in file.readlines():
                video_name = video_info.split()[0]
                filter_video_dict.add(video_name)
        print('Proc:{}, Filter File Count: {}'.format(id, len(filter_video_dict)))

        video_list = []
        with open(os.path.join(base_path, video_list_bbox_count_file.format(id)), 'r') as file:
            for video_info in file.readlines():
                video_info = video_info.split(',')
                # video_list.append([video_info[0].strip(), int(video_info[1].strip())])
                video_name = video_info[0].strip()
                video_list.append([video_name, int(video_info[1].strip())])
        print('Proc:{}, Read File BBox Count Finish!'.format(id))

        score_thr = THRESHOLDS[select_label]
        file_count = 0
        total_file_count = len(video_list)
        with open(os.path.join(base_path, video_list_bbox_file.format(id)), 'rb') as file:
            for video_name, BBox_count in video_list:
                if video_name not in filter_video_dict:
                    bbox_info = file.read(28*BBox_count)
                    if len(bbox_info) < 28*BBox_count:
                        print('Proc:{}, File Count: {}/{}'.format(id, file_count, total_file_count))
                        file_queue.put([None, True, id])
                        return
                    continue
                bbox_infos = []
                for j in range(BBox_count):
                    bbox_info = file.read(28)
                    if len(bbox_info) < 28:
                        print('Proc:{}, File Count: {}/{}'.format(id, file_count, total_file_count))
                        file_queue.put([None, True, id])
                        return
                    bbox = struct.unpack('6i1f', bbox_info)
                    label = bbox[1]
                    if select_label != label:
                        continue
                    # frame_id = bbox[0]
                    score = bbox[6]
                    if score < score_thr:
                        continue
                    bbox_infos.append(bbox_info)

                if len(bbox_infos) > 0:
                    file_queue.put([(video_name, bbox_infos), False, id])

                file_count += 1
                if file_count % 1000 == 0:
                    print('Proc:{}, File Count: {}/{}'.format(id, file_count, total_file_count))
                    # return
    except Exception as e:
        print('Proc:{}, File Count: {}/{}'.format(id, file_count, total_file_count))
    print('Proc:{}, File Count: {}/{}'.format(id, file_count, total_file_count))
    file_queue.put([None, True, id])


def main():
    if not os.path.exists(out_path):
        os.makedirs(out_path)
    os.system('chmod a+wr {}'.format(out_path))

    file_queue = multiprocessing.Queue()
    select_video_workers = []
    for i in range(PROCESS_COUNT):
        select_video_workers.append(multiprocessing.Process(target=filter_proc, args=(file_queue, i,)))

    for i in range(PROCESS_COUNT):
        select_video_workers[i].start()

    file_path = os.path.join(out_path, '{}_video_list_bbox_count.txt'.format(select_class))
    bbox_path = os.path.join(out_path, '{}_video_list_bbox.dat'.format(select_class))

    file = open(file_path, 'w')
    bbox_file = open(bbox_path, 'wb')

    file_count = 0
    try:
        finish_worker_count = 0
        while True:
            file_info = file_queue.get(block=True)
            if file_info is None:
                break
            video_info, finish, id = file_info
            if finish:
                print('Proc{} finish'.format(id, ))
                finish_worker_count += 1
                if PROCESS_COUNT <= finish_worker_count:
                    break
                continue
            if (video_info is None) or (len(video_info) < 1):
                continue
            video_name, bbox_infos = video_info
            bbox_count = len(bbox_infos)
            if bbox_count > 0:
                file_info_str = '{},{}\n'.format(video_name, bbox_count)
                file.write(file_info_str)
                for bbox_info in bbox_infos:
                    bbox_file.write(bbox_info)
            file_count += 1
            if file_count % 100 == 0:
                print('Process File Count:', file_count)
    except Exception as e:
        print(e)

    file.close()
    bbox_file.close()
    os.system('chmod a+wr {}'.format(file_path))
    os.system('chmod a+wr {}'.format(bbox_path))

    for i in range(PROCESS_COUNT):
        select_video_workers[i].join()

    print('finish!!!!')


if __name__ == '__main__':
    main()
